package com.csust.automaticteach.demo;

import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.Base64;

public class Embedding {

    static class AssembleHeaderException extends Exception {
        public AssembleHeaderException(String msg) {
            super(msg);
        }
    }

    static class Url {
        String host;
        String path;
        String schema;

        public Url(String host, String path, String schema) {
            this.host = host;
            this.path = path;
            this.schema = schema;
        }
    }

    private static String sha256base64(byte[] data) throws Exception {
        MessageDigest sha256 = MessageDigest.getInstance("SHA-256");
        byte[] digest = sha256.digest(data);
        return Base64.getEncoder().encodeToString(digest);
    }

    private static Url parseUrl(String requestUrl) throws AssembleHeaderException {
        int stidx = requestUrl.indexOf("://");
        if (stidx == -1) {
            throw new AssembleHeaderException("invalid request url: " + requestUrl);
        }

        String schema = requestUrl.substring(0, stidx + 3);
        String hostPart = requestUrl.substring(stidx + 3);

        int edidx = hostPart.indexOf("/");
        if (edidx <= 0) {
            throw new AssembleHeaderException("invalid request url: " + requestUrl);
        }

        String host = hostPart.substring(0, edidx);
        String path = hostPart.substring(edidx);

        return new Url(host, path, schema);
    }

    private static String assembleWsAuthUrl(String requestUrl, String method, String apiKey, String apiSecret)
            throws Exception {
        Url u = parseUrl(requestUrl);
        String host = u.host;
        String path = u.path;

        Calendar calendar = Calendar.getInstance();
        SimpleDateFormat sdf = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
        sdf.setTimeZone(TimeZone.getTimeZone("GMT"));
        String date = sdf.format(calendar.getTime());

        String signatureOrigin = "host: " + host + "\ndate: " + date + "\n" + method + " " + path + " HTTP/1.1";
        Mac mac = Mac.getInstance("HmacSHA256");
        SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "HmacSHA256");
        mac.init(spec);
        byte[] signatureSha = mac.doFinal(signatureOrigin.getBytes(StandardCharsets.UTF_8));
        String signature = Base64.getEncoder().encodeToString(signatureSha);

        String authorizationOrigin = String.format(
                "api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"",
                apiKey, "hmac-sha256", "host date request-line", signature);
        String authorization = Base64.getEncoder().encodeToString(authorizationOrigin.getBytes(StandardCharsets.UTF_8));

        Map<String, String> values = new LinkedHashMap<>();
        values.put("host", host);
        values.put("date", date);
        values.put("authorization", authorization);

        StringBuilder queryString = new StringBuilder();
        for (Map.Entry<String, String> entry : values.entrySet()) {
            if (queryString.length() > 0) {
                queryString.append("&");
            }
            queryString.append(URLEncoder.encode(entry.getKey(), "UTF-8"));
            queryString.append("=");
            queryString.append(URLEncoder.encode(entry.getValue(), "UTF-8"));
        }

        return requestUrl + "?" + queryString.toString();
    }

    private static JSONObject getBody(String appid, String text, String style) {
        String encodedText = Base64.getEncoder().encodeToString(text.getBytes(StandardCharsets.UTF_8));

        JSONObject body = new JSONObject();
        JSONObject header = new JSONObject();
        header.put("app_id", appid);
        header.put("uid", "39769795890");
        header.put("status", 3);
        body.put("header", header);

        JSONObject parameter = new JSONObject();
        JSONObject emb = new JSONObject();
        emb.put("domain", style);

        JSONObject feature = new JSONObject();
        feature.put("encoding", "utf8");
        emb.put("feature", feature);
        parameter.put("emb", emb);
        body.put("parameter", parameter);

        JSONObject payload = new JSONObject();
        JSONObject messages = new JSONObject();
        messages.put("text", encodedText);
        payload.put("messages", messages);
        body.put("payload", payload);

        return body;
    }

    private static String getEmbedding(String text, String appid, String apikey, String apisecret, String style)
            throws Exception {
        String host = "https://emb-cn-huabei-1.xf-yun.com/";
        String url = assembleWsAuthUrl(host, "POST", apikey, apisecret);

        JSONObject body = getBody(appid, text, style);

        URL obj = new URL(url);
        HttpURLConnection con = (HttpURLConnection) obj.openConnection();
        con.setRequestMethod("POST");
        con.setRequestProperty("Content-Type", "application/json");
        con.setDoOutput(true);

        try (OutputStream os = con.getOutputStream()) {
            byte[] input = body.toString().getBytes(StandardCharsets.UTF_8);
            os.write(input, 0, input.length);
        }

        StringBuilder response = new StringBuilder();
        try (BufferedReader br = new BufferedReader(
                new InputStreamReader(con.getInputStream(), StandardCharsets.UTF_8))) {
            String responseLine;
            while ((responseLine = br.readLine()) != null) {
                response.append(responseLine.trim());
            }
        }
        return response.toString();
    }

    public static String getEmbqEmbedding(String text, String appid, String apikey, String apisecret)
            throws Exception {
        return getEmbedding(text, appid, apikey, apisecret, "query");
    }

    public static String getEmbpEmbedding(String text, String appid, String apikey, String apisecret)
            throws Exception {
        return getEmbedding(text, appid, apikey, apisecret, "para");
    }

    private static float[] parseMessage(String message) {
        JSONObject data = new JSONObject(message);
        int code = data.getJSONObject("header").getInt("code");
        if (code != 0) {
            System.err.println("请求错误: " + code + ", " + data);
            return null;
        }

//        String sid = data.getJSONObject("header").getString("sid");
        String sid = "";
        String textBase = "";
        System.out.println("本次会话的id为：" + sid);

//        String textBase = data.getJSONObject("payload")
//                .getJSONObject("feature")
//                .getString("text");
        byte[] textData = Base64.getDecoder().decode(textBase);

        ByteBuffer buffer = ByteBuffer.wrap(textData);
        buffer.order(ByteOrder.LITTLE_ENDIAN);
        float[] vector = new float[textData.length / 4];
        buffer.asFloatBuffer().get(vector);

        System.out.println("向量维度: " + vector.length);
        System.out.println("返回的向量化数组为:");
        System.out.println(Arrays.toString(vector));

        return vector;
    }

    public static void main(String[] args) {
        try {
            String APPID = "XXXXXXXX";
            String APISecret = "XXXXXXXXXXXXXXXXXXXXXXXX";
            String APIKEY = "XXXXXXXXXXXXXXXXXXXXXXXX";

            JSONObject desc = new JSONObject();
            JSONArray messages = new JSONArray();
            JSONObject msg = new JSONObject();
            msg.put("content", "这段话的内容变成向量化是什么样的");
            msg.put("role", "user");
            messages.put(msg);
            desc.put("messages", messages);

            // 知识库内容向量化
            String res = getEmbpEmbedding(desc.toString(), APPID, APIKEY, APISecret);
            float[] vector = parseMessage(res);

            // 用户问题向量化 (示例)
            // String ques = "{\"messages\":[{\"content\":\"这段话的内容变成向量化是什么样的\",\"role\":\"user\"}]}";
            // res = getEmbqEmbedding(ques, APPID, APIKEY, APISecret);
            // vector = parseMessage(res);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}